from django.db.backends.postgresql.base import *
from django.db.backends.postgresql.introspection import DatabaseIntrospection

class OpenGaussDatabaseIntrospection(DatabaseIntrospection):
    def get_constraints(self, cursor, table_name):
        """
        重写get_constraints方法，避免使用WITH ORDINALITY语法
        这是一个简化版本，可能需要根据实际情况调整
        """
        constraints = {}
        # 获取表的主键
        cursor.execute("""
            SELECT kc.column_name
            FROM information_schema.table_constraints tc
            JOIN information_schema.key_column_usage kc ON
                kc.constraint_name = tc.constraint_name AND
                kc.table_schema = tc.table_schema AND
                kc.table_name = tc.table_name
            WHERE tc.constraint_type = 'PRIMARY KEY' AND tc.table_name = %s
        """, [table_name])
        primary_key_columns = [row[0] for row in cursor.fetchall()]
        
        if primary_key_columns:
            constraints['pk'] = {
                'columns': primary_key_columns,
                'primary_key': True,
                'unique': True,
                'foreign_key': None,
                'check': False,
                'index': True,
            }
        
        # 获取索引
        cursor.execute("""
            SELECT
                i.relname as index_name,
                array_agg(a.attname ORDER BY c.ordinality) as column_names,
                ix.indisunique as unique,
                ix.indisprimary as primary
            FROM
                pg_catalog.pg_class i,
                pg_catalog.pg_class t,
                pg_catalog.pg_index ix,
                pg_catalog.pg_attribute a,
                generate_series(1, current_setting('max_index_keys')::int) c(ordinality)
            WHERE
                i.oid = ix.indexrelid AND
                t.oid = ix.indrelid AND
                a.attrelid = t.oid AND
                t.relname = %s AND
                a.attnum = ix.indkey[c.ordinality-1] AND
                a.attnum > 0
            GROUP BY
                i.relname, ix.indisunique, ix.indisprimary
        """, [table_name])
        
        for index_name, column_names, unique, primary in cursor.fetchall():
            if primary:
                continue  # 已经处理过主键
            constraints[index_name] = {
                'columns': column_names,
                'primary_key': primary,
                'unique': unique,
                'foreign_key': None,
                'check': False,
                'index': True,
            }
        
        # 获取外键
        cursor.execute("""
            SELECT
                tc.constraint_name,
                kcu.column_name,
                ccu.table_name AS referenced_table,
                ccu.column_name AS referenced_column
            FROM information_schema.table_constraints AS tc
            JOIN information_schema.key_column_usage AS kcu
              ON tc.constraint_name = kcu.constraint_name
            JOIN information_schema.constraint_column_usage AS ccu
              ON ccu.constraint_name = tc.constraint_name
            WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name = %s
        """, [table_name])
        
        for constraint_name, column_name, referenced_table, referenced_column in cursor.fetchall():
            constraints[constraint_name] = {
                'columns': [column_name],
                'primary_key': False,
                'unique': False,
                'foreign_key': (referenced_table, referenced_column),
                'check': False,
                'index': False,
            }
        
        return constraints

class DatabaseWrapper(DatabaseWrapper):
    introspection_class = OpenGaussDatabaseIntrospection
